import os
from openai import OpenAI

class ImagePromptGenerator:
    def __init__(self, config=None):
        """初始化插画提示词生成器"""
        self.config = config or {}
        # 创建OpenAI客户端，使用与故事生成器相同的配置
        self.client = self.create_client()
    
    def create_client(self):
        """创建OpenAI客户端"""
        try:
            # 硬编码的API密钥
            api_key = "b8fa9da4-5c3a-4617-b17e-ff0cb092627f"
            
            client = OpenAI(
                # 配置API端点
                base_url="https://ark.cn-beijing.volces.com/api/v3",
                api_key=api_key,
            )
            
            return client
        except Exception as e:
            print(f"初始化OpenAI客户端失败: {e}")
            # 即使初始化失败，也让程序继续运行，后面的方法会处理错误
            return None
    
    def generate_prompts(self, story_content):
        """从故事内容生成插画提示词"""
        if not self.client:
            # 如果没有有效的API客户端，返回示例提示词
            return self._get_sample_prompts(story_content)
        
        try:
            # 构建提示词
            prompt = self._build_prompt_generation_prompt(story_content)
            
            # 调用OpenAI API生成插画提示词
            response = self.client.chat.completions.create(
                # 使用与故事生成器相同的模型
                model="ep-20251001204558-gqnb8",
                messages=[
                    {"role": "system", "content": "你是一位专业的插画师，擅长为儿童故事创作插画提示词。"},
                    {"role": "user", "content": prompt}
                ],
                max_tokens=1000,
                temperature=0.8
            )
            
            # 提取生成的提示词列表
            prompts_text = response.choices[0].message.content.strip()
            
            # 解析提示词列表
            prompts = self._parse_prompts(prompts_text)
            
            return prompts
        except Exception as e:
            print(f"生成插画提示词失败: {e}")
            # 发生错误时，返回示例提示词
            return self._get_sample_prompts(story_content)
    
    def _build_prompt_generation_prompt(self, story_content):
        """构建生成插画提示词的提示词"""
        prompt = f"请为以下儿童故事生成{self.config['image_settings']['number_of_images']}个插画提示词，要求如下：\n"
        prompt += f"- 故事内容：{story_content}\n"
        prompt += f"- 插画风格：{self.config['image_settings']['style']}\n"
        prompt += "- 每个提示词应详细描述一个关键场景\n"
        prompt += "- 语言风格：生动、具体、适合AI绘画工具理解\n"
        prompt += "- 请包含角色、场景、动作和情感等元素\n"
        prompt += "- 请以列表形式返回，每个提示词单独一行\n"
        prompt += "- 不要包含额外的解释性文字"
        
        return prompt
    
    def _parse_prompts(self, prompts_text):
        """解析生成的提示词列表"""
        # 分割文本为行
        lines = prompts_text.strip().split('\n')
        
        # 过滤空行和可能的编号
        prompts = []
        for line in lines:
            line = line.strip()
            if line:
                # 移除可能的编号前缀（如"1. "、"- "等）
                if line[0].isdigit() and '.' in line[:3]:
                    prompt = line.split('.', 1)[1].strip()
                elif line.startswith('- '):
                    prompt = line[2:].strip()
                else:
                    prompt = line
                prompts.append(prompt)
        
        # 确保提示词数量符合要求
        if len(prompts) < self.config['image_settings']['number_of_images']:
            # 如果提示词数量不足，复制一些提示词
            while len(prompts) < self.config['image_settings']['number_of_images']:
                prompts.append(prompts[0])
        elif len(prompts) > self.config['image_settings']['number_of_images']:
            # 如果提示词数量过多，只取前几个
            prompts = prompts[:self.config['image_settings']['number_of_images']]
        
        return prompts
    
    def _get_sample_prompts(self, story_content):
        """当API调用失败时，返回示例提示词"""
        # 提取故事中的关键词作为示例提示词的基础
        keywords = self._extract_keywords(story_content)
        
        style = self.config['image_settings']['style']
        number_of_images = self.config['image_settings']['number_of_images']
        
        # 示例提示词模板
        sample_templates = [
            f"{{style}}风格的插画，可爱的{{character}}在{{setting}}里开心地玩耍，阳光明媚，色彩鲜艳",
            f"{{style}}风格的插画，{{character}}帮助朋友的温馨场景，{{setting}}背景，充满友谊的氛围",
            f"{{style}}风格的插画，{{character}}遇到困难时的表情，{{setting}}环境，生动有趣",
            f"{{style}}风格的插画，{{character}}解决问题后的快乐场景，{{setting}}背景，色彩丰富"
        ]
        
        # 生成示例提示词
        sample_prompts = []
        for i in range(number_of_images):
            template = sample_templates[i % len(sample_templates)]
            
            # 填充模板中的变量
            prompt = template.format(
                style=style,
                character=keywords.get('character', '小兔子'),
                setting=keywords.get('setting', '森林')
            )
            
            sample_prompts.append(prompt)
        
        return sample_prompts
    
    def _extract_keywords(self, story_content):
        """从故事内容中提取关键词"""
        keywords = {
            'character': '小兔子',
            'setting': '森林'
        }
        
        # 简单的关键词提取逻辑
        # 在实际应用中，可以使用更复杂的NLP方法
        if '小兔子' in story_content:
            keywords['character'] = '小兔子'
        elif '小熊' in story_content:
            keywords['character'] = '小熊'
        elif '小猫' in story_content:
            keywords['character'] = '小猫'
        elif '小狗' in story_content:
            keywords['character'] = '小狗'
        
        if '森林' in story_content:
            keywords['setting'] = '森林'
        elif '海洋' in story_content:
            keywords['setting'] = '海洋'
        elif '太空' in story_content:
            keywords['setting'] = '太空'
        elif '草原' in story_content:
            keywords['setting'] = '草原'
        
        return keywords
    
    def save_prompts(self, prompts, filename="image_prompts.txt"):
        """保存生成的插画提示词到文件"""
        with open(filename, "w", encoding="utf-8") as f:
            for i, prompt in enumerate(prompts, 1):
                f.write(f"{i}. {prompt}\n")
        return filename